import argparse
from typing import Dict, Any

import numpy as np
import ray
import json

import yaml

from modules import BacktrackingPredictor
from utils.logger import get_logger
from utils.tools import notify
from ray.data import DataContext

DataContext.get_current().enable_progress_bars = False

logger = get_logger()

parser = argparse.ArgumentParser(
    prog='backtraking',
    description='backtraking')

parser.add_argument("--inputs",
                    "-i",
                    type=str,
                    help="json argument")
parser.add_argument("--config_path",
                    "-c",
                    type=str,
                    help="json argument")
args = parser.parse_args()

with open(args.config_path) as fp:
    config = yaml.safe_load(fp)

request_data = json.loads(args.inputs)
callback_url = request_data.get("callback_url")
task_id = request_data.get("task_id")
tasks = request_data.get("tasks")
storage_cfg = config["storage"]
# mongo_client = MongoClient(storage_cfg['mongodb_uri'])
# db = mongo_client.get_database()
# collection = db['backtracking']


def set_progress_status(row: Dict[str, Any]) -> Dict[str, Any]:
    row["task_id"] = task_id
    row["callback_url"] = callback_url
    row["backtracking_progress"] = row.get("backtracking_progress", "pending")
    return row


# notify(tasks, callback_url, task_id, "pending")

ray_cfg = config['model']['ray_options']

results = []
for row in (
        ray.data.from_items(tasks, override_num_blocks=ray_cfg['override_num_blocks'])
                .map(set_progress_status)
                .map(BacktrackingPredictor,
                     concurrency=ray_cfg['concurrency'],
                     num_cpus=ray_cfg['num_cpus'],
                     fn_constructor_kwargs={"cfg": config},
                     )
                .iter_rows()
):
    # del row['silent_points']
    # notify([row], callback_url, task_id, "part_finished")
    results.append(row)

# notify(results, callback_url, task_id, "all_finished")
# documents = collection.insert_many(results)

# 输出结果
# print(f"Inserted documents with IDs: {documents.inserted_ids}")
# print("Documents inserted successfully.")
